import warnings
import faiss
import os
import time
import numpy as np
import torch
from torchvision import models
from torchvision import transforms
from torch.utils.data import DataLoader
from PIL import Image
import matplotlib.pyplot as plt
from model import transform, load_model, feature_extract
from skimage import io
from config import *


def index_search(index, feat, topK):
    """定义检索方法
    :param feat: 检索的图片特征
    :param topK: 返回最高topK相似的图片
    :return:
    """

    feat = np.expand_dims(np.array(feat), axis=0)  # 将feat数组进行扩展，增加一维
    feat = feat.astype('float32')  # 数据类型转换为float32

    start_time = time.time()
    dis, ind = index.search(feat, topK)  # 递归调用
    end_time = time.time()

    print('index_search consume time:{}ms'.format(int(end_time - start_time) * 1000))  # 打印查询花费的时间
    return dis, ind  # 距离，相似图片id


def visual_plot(img_folder, filename, ind, dis, topK, query_img:Image=None):
    """

    :param ind: 搜索结果的索引
    :param dis: 搜索结果的距离
    :param topK: 返回的最相似的前k个结果
    :param query_img: （可选）作为查询的图像，用于在可视化中显示
    :return:
    """

    # 相似照片
    cols = 4  # 每一行显示的照片数
    rows = int(topK / cols)  # 总共需要的行数
    idx = 0  # idx初始化为0

    fig, axes = plt.subplots(rows, cols, figsize=(20, 5 * rows), tight_layout=True)
    """
    rows 和 cols 是子图的行数和列数，分别通过 topK 变量计算得出。
    figsize 是整个图形对象的大小，设置为 (20, 5*rows)，表示宽度为 20，高度为每行 5 个子图的高度总和。
    tight_layout 表示是否自动调整子图之间的间距，使得整个图形更美观。
    """
    # axes[0,0].imshow(query_img)

    for row in range(rows):
        for col in range(cols):
            _id = ind[0][idx]
            _dis = dis[0][idx]

            if _id == -1:
                continue

            img_path = os.path.join(img_folder, '{}.jpg'.format(_id))
            # print(img_path)

            if query_img is not None and idx == 0:
                axes[row, col].imshow(query_img)
                axes[row, col].set_title('query', fontsize=20)
            else:
                img = plt.imread(img_path)
                axes[row, col].imshow(img)
                axes[row, col].set_title('matched_-{}_{}'.format(_id, int(_dis)), fontsize=20)
            idx += 1

    folder_name = res_folder

    # 保存结果图片
    plt.savefig(os.path.join(folder_name, filename.rsplit('.', 1)[0] + "_res"))

def search(image, filename):

    # ============ 拿到索引 ==============
    index = faiss.read_index(index_file_name)
    """拿到index_file.index"""

    # ============ 拿到设备 ==============
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    # ============ 拿到文件目录 img_folder ==============
    data_root = os.path.abspath(os.path.join(os.getcwd(), ".."))  # 拿到 souTU
    # print(data_root)
    img_folder = os.path.join(data_root, img_folder_name)  # 拿到 souTu/VOC2012_small
    # print(image_path)
    assert os.path.exists(img_folder), "{} path does not exist.".format(img_folder)  # 断言警告

    # ============ 载入模型 ==============
    model = load_model()

    topK = 10  # 需要返回20张图片

    # 检测的图片名称
    # img_src = search_pic_name

    # img = io.imread(img_src)  # imread函数读取了一个图像文件，该函数返回一个numpy数组类型的图像

    img = image

    # img = Image.fromarray(img)  # ndarray_image为原来的numpy数组类型的输入，将numpy数组类型的图像转换为Image对象类型，使用了Pillow库中的Image模块
    query_img = img  # 查询图片，Image对象

    img = img.convert("RGB")  # 将图像转换为RGB格式
    img = transform(img)  # torch.Size([3, 224, 224])  #将图像转换为大小为3x224x224的张量
    img = img.unsqueeze(0)  # 使用了PyTorch的unsqueeze函数，将维度为1的新维度添加到张量的第一维，即将大小为[3, 224, 224]的张量变为大小为[1, 3, 224, 224]的张量
    img = img.to(device)  # 将张量移动到指定的设备（如CPU或GPU）上进行处理

    with torch.no_grad():  # (简单说就是不计算模型的梯度，因为梯度是为了求权重的)使用no_grad函数可以临时禁用梯度计算，以加快模型的推理速度。同时，这也可以避免在推理时意外修改模型的权重。使用no_grad函数可以临时禁用梯度计算，以加快模型的推理速度。同时，这也可以避免在推理时意外修改模型的权重。
        print("提取图片特征")
        feature = feature_extract(model, img)

        print("")
        # 使用了tolist函数将特征转换为Python列表格式，并取出其中的第一个元素。由于特征通常是一个大小为[1, n]的张量，
        # 因此使用[0]索引可以获取特征列表中的唯一元素，即大小为[n]的一维列表。
        feature_list = feature.data.cpu().tolist()[0]

        print("")

        dis, ind = index_search(index, feature_list, topK=topK)

        # print(dis)
        print(ind)  # 打印最相似的索引
        print('4.****************')
        visual_plot(img_folder, filename, ind, dis, topK, query_img)

# if __name__ == '__main__':
#     main()